from pathlib import Path
from triple_graph import TripleGraph
from utils import read_triples, read_mapping, read_rules
from os import path as os_path
from os import makedirs


class Graph_Completion(object):
    def __init__(self, directory, rule_path):
        self.triple_graph = TripleGraph()
        self.load_data(directory, rule_path)
        self.triple_graph.load(self.triples)


    def load_data(self, directory, rule_path):
        self.entity2id = read_mapping(directory / 'entity2id.txt')           # {str: int}
        self.relation2id = read_mapping(directory / 'relation2id.txt')       # {str: int}
        self.triples = read_triples(directory / 'triples.txt')               # [(int(head), int(tail), int(relation)), ...]
        self.rules = read_rules(rule_path / 'parsed_rules.txt', self.relation2id)     # [(((str, str, int(relation)), (str, str, int),...) #premises, (str, str, int) #hypothesis, conf), ...]
        self.id2entity = {i: entity for entity, i in self.entity2id.items()}
        self.id2relation = {i: relation for relation, i in self.relation2id.items()}


    def _completion(self):
        triples = self.triples
        new_triple_conf_dict, new_triple_premises_dict = {}, {}
        for rule in self.rules:
            new_triple_conf_premises_candidates = self.triple_graph.inference_by_rule(rule)

            for triple, conf, premises in new_triple_conf_premises_candidates:
                if triple not in triples:
                    if triple not in new_triple_conf_dict:
                        new_triple_conf_dict[triple] = conf
                        new_triple_premises_dict[triple] = premises
                    else:
                        if new_triple_conf_dict[triple] < conf:
                            new_triple_conf_dict[triple] = conf
                            new_triple_premises_dict[triple] = premises
        return new_triple_conf_dict, new_triple_premises_dict

    def completion(self, result_directory):
        print("Start completion ...")
        if not os_path.exists(result_directory):
            makedirs(result_directory)
        print("Start instantiate ...")
        triple_conf_dict, triple_premises_dict = self._completion()
        print("End instantiate. Start write to file")
        assert len(triple_conf_dict) == len(triple_premises_dict)
        new_triples, probabilities = [], []
        with open(os_path.join(result_directory, "new_triple_conf.txt"), "w", encoding="utf8") as fout:
            for triple in triple_conf_dict:
                triple_str = [self.id2entity[triple[0]], self.id2relation[triple[2]], self.id2entity[triple[1]]] \
                             + [str(triple_conf_dict[triple])]
                fout.write("\t".join(triple_str) + "\n")
                new_triples.append('\t'.join(triple_str[:3]))
                probabilities.append(triple_str[3])

        with open(os_path.join(result_directory, "new_triple_premises.txt"), "w", encoding="utf8") as fout:
            for triple in triple_premises_dict:
                triple_str = "\t".join([self.id2entity[triple[0]], self.id2relation[triple[2]], self.id2entity[triple[1]]])
                premises_str = []
                for premises in triple_premises_dict[triple]:
                    premises_str.append("[" + self.id2entity[premises[0]] + "\t" + self.id2relation[premises[2]]
                                        + "\t" + self.id2entity[premises[1]] + "]")
                premises_str = "\t".join(premises_str)
                fout.write("\t<=\t".join([triple_str, premises_str]) + "\n")

        with open(os_path.join(result_directory, "new_triples.txt"), "w", encoding="utf8") as fout:
            fout.write(str(len(triple_conf_dict)) + "\n")
            triples = triple_conf_dict.keys()
            triples = ["\t".join([self.id2entity[triple[0]], self.id2relation[triple[2]], self.id2entity[triple[1]]]) for triple in triples]
            for triple in triples:
                fout.write(triple + "\n")

        return new_triples, probabilities


if __name__ == "__main__":
    data_path = os_path.join("..", "data", "DBP15k")
    directory = Path(os_path.join(data_path))
    result_directory = os_path.join("..", "result", "DBP15k")
    if not os_path.join(result_directory):
        makedirs(result_directory)
    rule_path = Path(os_path.join('..', 'result', 'DBP15k'))
    GC = Graph_Completion(directory, rule_path)
    # print("end")
    triples, probs = GC.completion(result_directory)